#!/bin/bash

# Runs the "345M" parameter model
. path.sh
export CUDA_DEVICE_MAX_CONNECTIONS=1
export NCCL_LAUNCH_TIMEOUT=30
#export NCCL_SOCKET_IFNAME=eth1
#export NCCL_IB_GID_INDEX=3
#export NCCL_IB_HCA=mlx5_2:1
#export NCCL_IB_SL=3
#export NCCL_CHECK_DISABLE=1
export NCCL_P2P_DISABLE=0
#export NCCL_IB_DISABLE=1
#export NCCL_DEBUG=INFO
export NCCL_LL_THRESHOLD=16384
#export NCCL_IB_CUDA_SUPPORT=1

# export NVTE_APPLY_QK_LAYER_SCALING=1

# DATA_PATH=data/data_v18.16.train.jsonl.${SEQ_LENGTH}.hunyuan.json.ep1.glm3.sample_merge.pt
#VALID_DATA_PATH=data/data_v18.16.valid.jsonl.${SEQ_LENGTH}.hunyuan.json.ep1.glm3.no_sample_merge.pt
seq_len=8192
data_version=19.12
VALID_DATA_PATH=workspace_sh/data_v${data_version}.valid.jsonl.${seq_len}.hunyuan.json.ep1.qwen2.sample_merge.pt

DATA_PATH=workspace_sh/data_v${data_version}.train.jsonl.${seq_len}.hunyuan.json.ep1.qwen2.sample_merge.pt

GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=localhost
MASTER_PORT=12323
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

SEQ_LENGTH=$seq_len
SAVE_PREFIX=qwen2.5_mglm_debug
. utils/parse_options.sh



MICRO_BATCH_SIZE=1
GLOBAL_BATCH_SIZE=64

TP=4
PP=2
CP=1
# require to align with weight dimensions
HIDDEN_SIZE=3584
FFN_HIDDEN_SIZE=18944
NUM_LAYERS=28
NUM_HEADS=28
VOCAB_SIZE=152064
######################################
SAVE_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/exp/exp_meetingGPT/${SAVE_PREFIX}_TP${TP}_PP${PP}_CP${CP}_MBZ${MICRO_BATCH_SIZE}_GBSZ${GLOBAL_BATCH_SIZE}_seq${SEQ_LENGTH}
#SAVE_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/exp/exp_meetingGPT/qwen2.5_mglm_datav18.16_TP${TP}_PP${PP}_CP${CP}_MBZ${MICRO_BATCH_SIZE}_GBSZ${GLOBAL_BATCH_SIZE}_seq${SEQ_LENGTH}
#SAVE_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/exp/exp_meetingGPT/debug


HF_LLAMA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/Qwen/Qwen2.5-7B


# TOKENIZER_MODEL=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/Llama-2-7b-chat-hf


TOKENIZER_TYPE=GPT2BPETokenizer


merge_path=$HF_LLAMA_PATH/merges.txt
vocab_path=$HF_LLAMA_PATH/vocab.json

# CHECKPOINT_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/Megatron-LM/hunyuan-7b-mega-T2P2

#CHECKPOINT_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/Megatron-LM/llama7b-megatron-TP${TP}-PP${PP}
CHECKPOINT_PATH=workspace_sh/qwen2.5-7b-instruct-megatron-TP${TP}-PP${PP}-TE
tensorboard_output=$SAVE_PATH/tensorboard
mkdir -p $SAVE_PATH/tensorboard

if [ ! -d "$SAVE_PATH/iter_0000001" ] && [ "$NODE_RANK" -eq 0 ]; then
    echo "$SAVE_PATH/iter_0000001 not found, copying from $CHECKPOINT_PATH"
    cp -r $CHECKPOINT_PATH/* $SAVE_PATH/
fi

DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"

GPT_ARGS="
    --tensor-model-parallel-size $TP \
    --pipeline-model-parallel-size $PP \
    --context-parallel-size $CP \
    --sequence-parallel \
    --num-layers $NUM_LAYERS \
    --hidden-size $HIDDEN_SIZE \
    --ffn-hidden-size $FFN_HIDDEN_SIZE \
    --num-attention-heads $NUM_HEADS \
    --seq-length $SEQ_LENGTH \
    --max-position-embeddings $SEQ_LENGTH \
    --micro-batch-size $MICRO_BATCH_SIZE \
    --global-batch-size $GLOBAL_BATCH_SIZE \
    --lr 1e-5 \
    --min-lr 1e-6 \
    --train-iters 600 \
    --lr-decay-iters 550 \
    --lr-decay-style cosine \
    --weight-decay 1e-1 \
    --attention-dropout 0 \
    --hidden-dropout 0 \
    --lr-warmup-iters 0 \
    --clip-grad 1.0 \
    --recompute-activations \
    --recompute-granularity selective \
    --use-distributed-optimizer \
    --optimizer hybridadam \
    --exit-on-missing-checkpoint \
    --use-rotary-position-embeddings \
    --normalization RMSNorm \
    --no-position-embedding \
    --disable-bias-linear \
    --swiglu \
    --no-masked-softmax-fusion \
    --transformer-impl transformer_engine \
    --ckpt-format torch \
    --dataloader-type cyclic \
    --tokenizer-type $TOKENIZER_TYPE \
    --vocab-file $vocab_path \
    --merge-file $merge_path \
    --untie-embeddings-and-output-weights \
    --attention-softmax-in-fp32 \
    --num-query-groups 4 \
    --make-vocab-size-divisible-by 38016 \
    --group-query-attention \
    --norm-epsilon 1e-6 \
    --rotary-base 1000000 \
    --add-qkv-bias \
    --reset-position-ids \
    --reset-attention-mask \
    --tensorboard-dir $tensorboard_output \
    --bf16
"


    #--use-checkpoint-args \

    #--use-legacy-models \

#    --tokenizer-type Llama2Tokenizer \
 #   --tokenizer-model ${TOKENIZER_MODEL} \

    #--untie-embeddings-and-output-weights \
    # --vocab-size 290943 \
    # --mock-data \
    # --data-path $DATA_PATH \
DATA_ARGS="
    --valid-data-path $VALID_DATA_PATH \
    --train-data-path $DATA_PATH
"

OUTPUT_ARGS="
    --log-interval 1 \
    --save-interval 50 \
    --eval-interval 50 \
    --log-throughput \
    --log-memory-to-tensorboard \
    --log-timers-to-tensorboard \
    --no-load-optim \
    --no-load-rng \
    --eval-iters 16
"

torchrun $DISTRIBUTED_ARGS pretrain_qwen.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    --distributed-backend nccl \
    --save $SAVE_PATH \
    --load $CHECKPOINT_PATH



